import os
import shutil


def copy_matching_images(txt_dir, jpg_dir, dest_dir):
    """
    将匹配的txt文件和jpg图片复制到目标文件夹

    参数:
    txt_dir: 包含txt文件的目录
    jpg_dir: 包含jpg图片的目录
    dest_dir: 目标输出目录
    """
    # 确保目标目录存在
    os.makedirs(dest_dir, exist_ok=True)

    # 获取txt目录中所有不带扩展名的文件名
    txt_files = [os.path.splitext(f)[0] for f in os.listdir(txt_dir) if f.endswith('.txt')]

    # 计数器
    copied_count = 0
    missing_count = 0

    print(f"在 {txt_dir} 中找到 {len(txt_files)} 个txt文件")
    print(f"正在匹配 {jpg_dir} 中的jpg图片...")

    # 遍历所有txt文件名
    for base_name in txt_files:
        jpg_path = os.path.join(jpg_dir, base_name + '.jpg')

        # 检查jpg文件是否存在
        if os.path.isfile(jpg_path):
            dest_path = os.path.join(dest_dir, base_name + '.jpg')
            shutil.copy2(jpg_path, dest_path)
            copied_count += 1
        else:
            print(f"警告: 未找到匹配的图片: {base_name}.jpg")
            missing_count += 1

    # 输出结果
    print("\n操作完成!")
    print(f"成功复制 {copied_count} 个jpg图片到 {dest_dir}")
    print(f"未找到匹配的图片: {missing_count} 个")

    return copied_count


if __name__ == "__main__":
    # 输入目录配置 - 修改为你的实际路径
    txt_directory = r"K:\1-YolotrainPic\20250624\xml_outpu2\yolo_output2"  # 包含txt文件的目录
    jpg_directory = r"K:\1-YolotrainPic\20250624\images"  # 包含jpg图片的目录
    destination_directory = r"K:\1-YolotrainPic\20250624\xml_outpu2\yolo_images2"  # 目标输出目录

    # 执行复制操作
    copy_matching_images(txt_directory, jpg_directory, destination_directory)